#include "stdafx.h"
#include "net/HTTPDigestCredentials.hpp"
#include "net/HTTPRequest.hpp"
#include "net/HTTPResponse.hpp"

#include "code/digestengine.hpp"
#include "code/md5engine.hpp"
using JHCPP::code::CMD5Engine;
using JHCPP::code::CDigestEngine;
#include "utils/numberformatter.hpp"
using JHCPP::utils::CNumberFormatter;
#include "stdlib/timestamp.hpp"
#include "stdlib/exception.hpp"
using JHCPP::stdlib::CTimeStamp;
using namespace JHCPP::stdlib::jhException;
#include "stream/stringtokenizer.hpp"
using JHCPP::stream::CStringTokenizer;


NET_NAMESPACE_BEGIN

std::string digest(CDigestEngine& engine,
				   const std::string& a,
				   const std::string& b,
				   const std::string& c = std::string(),
				   const std::string& d = std::string(),
				   const std::string& e = std::string(),
				   const std::string& f = std::string())
{
	engine.reset();
	engine.update(a);
	engine.update(':');
	engine.update(b);
	if (!c.empty()) 
	{
		engine.update(':');
		engine.update(c);
		if (!d.empty()) 
		{
			engine.update(':');
			engine.update(d);
			engine.update(':');
			engine.update(e);
			engine.update(':');
			engine.update(f);
		}
	}
	return CDigestEngine::digestToHex(engine.digest());
}

std::string formatNonceCounter(int counter)
{
	return CNumberFormatter::formatHex(counter, 8);
}


const std::string CHTTPDigestCredentials::SCHEME = "Digest";
const std::string CHTTPDigestCredentials::DEFAULT_ALGORITHM("MD5");
const std::string CHTTPDigestCredentials::DEFAULT_QOP("");
const std::string CHTTPDigestCredentials::NONCE_PARAM("nonce");
const std::string CHTTPDigestCredentials::REALM_PARAM("realm");
const std::string CHTTPDigestCredentials::QOP_PARAM("qop");
const std::string CHTTPDigestCredentials::ALGORITHM_PARAM("algorithm");
const std::string CHTTPDigestCredentials::USERNAME_PARAM("username");
const std::string CHTTPDigestCredentials::OPAQUE_PARAM("opaque");
const std::string CHTTPDigestCredentials::URI_PARAM("uri");
const std::string CHTTPDigestCredentials::RESPONSE_PARAM("response");
const std::string CHTTPDigestCredentials::AUTH_PARAM("auth");
const std::string CHTTPDigestCredentials::CNONCE_PARAM("cnonce");
const std::string CHTTPDigestCredentials::NC_PARAM("nc");

int CHTTPDigestCredentials::m_nonceCounter(0);
CMutex CHTTPDigestCredentials::m_nonceMutex;

CHTTPDigestCredentials::CHTTPDigestCredentials()
{
}

CHTTPDigestCredentials::CHTTPDigestCredentials(const std::string& username, const std::string& password)
	: m_username(username),
	m_password(password)
{
}

CHTTPDigestCredentials::~CHTTPDigestCredentials()
{
}

void CHTTPDigestCredentials::reset()
{
	m_requestAuthParams.clear();
	m_nc.clear();
}

void CHTTPDigestCredentials::setUsername(const std::string& username)
{
	m_username = username;
}

void CHTTPDigestCredentials::setPassword(const std::string& password)
{
	m_password = password;
}

const std::string& CHTTPDigestCredentials::getUsername() const
{
	return m_username;
}

const std::string& CHTTPDigestCredentials::getPassword() const
{
	return m_password;
}

void CHTTPDigestCredentials::authenticate(CHTTPRequest& request, const CHTTPResponse& response)
{
	authenticate(request, CHTTPAuthenticationParams(response));
}

void CHTTPDigestCredentials::authenticate(CHTTPRequest& request, const CHTTPAuthenticationParams& responseAuthParams)
{
	createAuthParams(request, responseAuthParams);
	request.setCredentials(SCHEME, m_requestAuthParams.toString());
}

void CHTTPDigestCredentials::updateAuthInfo(CHTTPRequest& request)
{
	updateAuthParams(request);
	request.setCredentials(SCHEME, m_requestAuthParams.toString());
}

void CHTTPDigestCredentials::proxyAuthenticate(CHTTPRequest& request, const CHTTPResponse& response)
{
	proxyAuthenticate(request, CHTTPAuthenticationParams(response, CHTTPAuthenticationParams::PROXY_AUTHENTICATE));
}

void CHTTPDigestCredentials::proxyAuthenticate(CHTTPRequest& request, const CHTTPAuthenticationParams& responseAuthParams)
{
	createAuthParams(request, responseAuthParams);
	request.setProxyCredentials(SCHEME, m_requestAuthParams.toString());
}

void CHTTPDigestCredentials::updateProxyAuthInfo(CHTTPRequest& request)
{
	updateAuthParams(request);
	request.setProxyCredentials(SCHEME, m_requestAuthParams.toString());
}

std::string CHTTPDigestCredentials::createNonce()
{
	CMutex::ScopedLock lock(m_nonceMutex);

	CMD5Engine md5;
	CTimeStamp::TimeVal now = CTimeStamp().epochMicroseconds();

	md5.update(&m_nonceCounter, sizeof(m_nonceCounter));
	md5.update(&now, sizeof(now));

	++m_nonceCounter;

	return CDigestEngine::digestToHex(md5.digest());
}

void CHTTPDigestCredentials::createAuthParams(const CHTTPRequest& request, const CHTTPAuthenticationParams& responseAuthParams)
{
	if (!responseAuthParams.has(NONCE_PARAM) || !responseAuthParams.has(REALM_PARAM))
		throw InvalidArgumentException("Invalid HTTP authentication parameters");

	const std::string& algorithm = responseAuthParams.get(ALGORITHM_PARAM, DEFAULT_ALGORITHM);

	if (icompare(algorithm, DEFAULT_ALGORITHM) != 0) 
		throw NotImplementedException("Unsupported digest algorithm", algorithm);

	const std::string& nonce = responseAuthParams.get(NONCE_PARAM);
	const std::string& qop = responseAuthParams.get(QOP_PARAM, DEFAULT_QOP);
	const std::string& realm = responseAuthParams.getRealm();

	m_requestAuthParams.clear();
	m_requestAuthParams.set(USERNAME_PARAM, m_username);
	m_requestAuthParams.set(NONCE_PARAM, nonce);
	m_requestAuthParams.setRealm(realm);
	if (responseAuthParams.has(OPAQUE_PARAM)) 
	{
		m_requestAuthParams.set(OPAQUE_PARAM, responseAuthParams.get(OPAQUE_PARAM));
	}

	if (qop.empty())
	{
		updateAuthParams(request);
	} 
	else
	{
		CStringTokenizer tok(qop, ",", CStringTokenizer::TOK_TRIM);
		bool qopSupported = false;
		for (CStringTokenizer::Iterator it = tok.begin(); it != tok.end(); ++it)
		{
			if (icompare(*it, AUTH_PARAM) == 0)
			{
				qopSupported = true;
				m_requestAuthParams.set(CNONCE_PARAM, createNonce());
				m_requestAuthParams.set(QOP_PARAM, *it);
				updateAuthParams(request);
				break;
			}
		}
		if (!qopSupported) 
			throw NotImplementedException("Unsupported QoP requested", qop);
	} 
}

void CHTTPDigestCredentials::updateAuthParams(const CHTTPRequest& request)
{
	CMD5Engine engine;
	const std::string& qop = m_requestAuthParams.get(QOP_PARAM, DEFAULT_QOP);
	const std::string& realm = m_requestAuthParams.getRealm();
	const std::string& nonce = m_requestAuthParams.get(NONCE_PARAM);

	m_requestAuthParams.set(URI_PARAM, request.getURI());

	if (qop.empty())
	{
		const std::string ha1 = digest(engine, m_username, realm, m_password);
		const std::string ha2 = digest(engine, request.getMethod(), request.getURI());

		m_requestAuthParams.set(RESPONSE_PARAM, digest(engine, ha1, nonce, ha2));
	}
	else if (icompare(qop, AUTH_PARAM) == 0) 
	{
		const std::string& cnonce = m_requestAuthParams.get(CNONCE_PARAM);

		const std::string ha1 = digest(engine, m_username, realm, m_password);
		const std::string ha2 = digest(engine, request.getMethod(), request.getURI());
		const std::string nc = formatNonceCounter(updateNonceCounter(nonce));

		m_requestAuthParams.set(NC_PARAM, nc);
		m_requestAuthParams.set(RESPONSE_PARAM, digest(engine, ha1, nonce, nc, cnonce, qop, ha2));
	}
}

bool CHTTPDigestCredentials::verifyAuthInfo(const CHTTPRequest& request) const
{
	CHTTPAuthenticationParams params(request);
	return verifyAuthParams(request, params);
}

bool CHTTPDigestCredentials::verifyAuthParams(const CHTTPRequest& request, const CHTTPAuthenticationParams& params) const
{
	const std::string& nonce = params.get(NONCE_PARAM);
	const std::string& realm = params.getRealm();
	const std::string& qop   = params.get(QOP_PARAM, DEFAULT_QOP);
	std::string response;
	CMD5Engine engine;
	if (qop.empty())
	{
		const std::string ha1 = digest(engine, m_username, realm, m_password);
		const std::string ha2 = digest(engine, request.getMethod(), request.getURI());
		response = digest(engine, ha1, nonce, ha2);
	}
	else if (icompare(qop, AUTH_PARAM) == 0) 
	{
		const std::string& cnonce = params.get(CNONCE_PARAM);
		const std::string& nc = params.get(NC_PARAM);
		const std::string ha1 = digest(engine, m_username, realm, m_password);
		const std::string ha2 = digest(engine, request.getMethod(), request.getURI());
		response = digest(engine, ha1, nonce, nc, cnonce, qop, ha2);
	}
	return response == params.get(RESPONSE_PARAM);
}

int CHTTPDigestCredentials::updateNonceCounter(const std::string& nonce)
{
	NonceCounterMap::iterator iter = m_nc.find(nonce);

	if (iter == m_nc.end()) 
	{
		iter = m_nc.insert(NonceCounterMap::value_type(nonce, 0)).first;
	}
	iter->second++;

	return iter->second;
}

NET_NAMESPACE_END
